#include <stdlib.h>
#include <stdio.h>
#include "cuda_runtime.h"
#include "nccl.h"

#if 1
#define CUDACHECK(cmd) do {                         \
  cudaError_t err = cmd;                            \
  if (err != cudaSuccess) {                         \
    printf("Failed: Cuda error %s:%d '%s'\n",       \
        __FILE__,__LINE__,cudaGetErrorString(err)); \
    exit(EXIT_FAILURE);                             \
  }                                                 \
} while(0)


#define NCCLCHECK(cmd) do {                         \
  ncclResult_t res = cmd;                           \
  if (res != ncclSuccess) {                         \
    printf("Failed, NCCL error %s:%d '%s'\n",       \
        __FILE__,__LINE__,ncclGetErrorString(res)); \
    exit(EXIT_FAILURE);                             \
  }                                                 \
} while(0)
#endif
void show_h_A_value(int *h_A, int cnt)
{
    for (int i = 0; i < cnt; ++i) {
        printf("%d ", (int)h_A[i]);
    }
    printf("\n");
}
#define MIN_ARGC                        (1)
#define MAX_SUPPORT_DEV_NUM_COUNT       (4)
char sCmdGetGPUNumber[] = "nvidia-smi topo -m";
int aSupportDevNum[MAX_SUPPORT_DEV_NUM_COUNT] = {1, 2, 4, 8};
static void broadcast_root(int nDev, int root, int nCount, int size, int* sendbuff_H, int* recvbuff_H, int** sendbuff, int** recvbuff, cudaStream_t* stream, ncclComm_t *comms);

int main(int argc, char* argv[])
{
  ncclComm_t comms[4];

  //managing 4 devices
  int nDev = 1;
  int nCount = 8;
  int size = nCount * sizeof(int);
  int devs[4] = { 0, 1, 2, 3 };

  int* sendbuff_H = (int*)malloc(size);
  int* recvbuff_H = (int*)malloc(size);
  //allocating and initializing device buffers
  int** sendbuff = NULL;
  int** recvbuff = NULL;
  cudaStream_t* stream = NULL;

  printf("=====================start==============\n");
  printf("[%s +%d %s] ncclBroadcast API test example\n", __FILE__, __LINE__, __func__);

    if(1 == argc)
    {
        printf("[%s +%d %s] argc=%d !=%d so exit(0)\n", __FILE__, __LINE__, __func__, argc, MIN_ARGC);
        printf("[%s +%d %s] example ./nccl_send_recv.elf num is the number can used \n", __FILE__, __LINE__, __func__);
        printf("./nccl_send_recv.elf 1 \n");
        printf("./nccl_send_recv.elf 2 \n");
        printf("./nccl_send_recv.elf 4 \n");
        system(sCmdGetGPUNumber);
        exit(0);
    }
    else
    {
      int i = 0;
      nDev = atoi(argv[1]);
      printf("[%s +%d %s] input nDev=%d\n", __FILE__, __LINE__, __func__, nDev);
      printf("[%s +%d %s] MAX_SUPPORT_DEV_NUM_COUNT=%d\n", __FILE__, __LINE__, __func__, MAX_SUPPORT_DEV_NUM_COUNT);
      for(i = 0; i < MAX_SUPPORT_DEV_NUM_COUNT; i++)
      {
        if(nDev == aSupportDevNum[i])
        {
          break;
        }
      }
      if(MAX_SUPPORT_DEV_NUM_COUNT == i)
      {
          printf("[%s +%d %s] nDev=%d is not in aSupportDevNum[] so exit(0)\n", __FILE__, __LINE__, __func__, nDev);
          printf("[%s +%d %s] example ./nccl_send_recv.elf num is the number can used \n", __FILE__, __LINE__, __func__);
          printf("./nccl_send_recv.elf 1 \n");
          printf("./nccl_send_recv.elf 2 \n");
          printf("./nccl_send_recv.elf 4 \n");
          printf("./nccl_send_recv.elf 8 \n");
          system(sCmdGetGPUNumber);
          exit(0);
      }
    }
  printf("[%s +%d %s] nCount=%d size=%d\n", __FILE__, __LINE__, __func__, nCount, size);
  sendbuff = (int**)malloc(nDev * sizeof(int*));
  recvbuff = (int**)malloc(nDev * sizeof(int*));
  stream = (cudaStream_t*)malloc(sizeof(cudaStream_t)*nDev);
  
  for (int i = 0; i < nDev; ++i)
  {
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaMalloc((void**)sendbuff + i, size));
    CUDACHECK(cudaMalloc((void**)recvbuff + i, size));
    CUDACHECK(cudaMalloc(sendbuff + i, size));
    CUDACHECK(cudaMalloc(recvbuff + i, size));
    #if 0
    printf("[%s +%d %s] set 1 to sendbuff by cudaMemset\n", __FILE__, __LINE__, __func__);
    printf("[%s +%d %s] set 0 to recvbuff by cudaMemset\n", __FILE__, __LINE__, __func__);
    CUDACHECK(cudaMemset(sendbuff[i], 1, size));
    CUDACHECK(cudaMemset(recvbuff[i], 0, size));
    #else
    for (int j = 0; j < nCount; ++j)
    {
      *(sendbuff_H + j) = j + 10 * i;
      //printf("[%s +%d %s] [%d] *(sendbuff_H + j=%d)=%d\n", __FILE__, __LINE__, __func__, i, j, *(sendbuff_H + j));
    }   
    printf("[%s +%d %s] copy sendbuff_H to sendbuff by cudaMemcpy\n", __FILE__, __LINE__, __func__);
    CUDACHECK(cudaMemcpy(sendbuff[i], sendbuff_H, size, cudaMemcpyHostToDevice));

    printf("[%s +%d %s] copy recvbuff to recvbuff_H by cudaMemcpy\n", __FILE__, __LINE__, __func__);
    CUDACHECK(cudaMemcpy(recvbuff[i], recvbuff_H, size, cudaMemcpyHostToDevice));
    for (int j = 0; j < nCount; ++j)
    {
      //printf("[%s +%d %s] [%d] *(recvbuff_H + j=%d)=%d\n", __FILE__, __LINE__, __func__, i, j, *(recvbuff_H + j));
    }
    #endif  
    CUDACHECK(cudaStreamCreate(stream+i));
  }

  //initializing NCCL
  printf("[%s +%d %s] before call ncclCommInitAll\n", __FILE__, __LINE__, __func__);
  NCCLCHECK(ncclCommInitAll(comms, nDev, devs));
  printf("[%s +%d %s] after call ncclCommInitAll\n", __FILE__, __LINE__, __func__);

  //calling NCCL communication API. Group API is required when using
   //multiple devices per thread
  int root = 0;
  #if 1
  root = 0;
  broadcast_root(nDev, root, nCount, size, sendbuff_H, recvbuff_H, sendbuff, recvbuff, stream, comms);
  #else
  root = 1;
  broadcast_root(nDev, root, nCount, size, sendbuff_H, recvbuff_H, sendbuff, recvbuff, stream, comms);
  #endif
  //free device buffers
  for (int i = 0; i < nDev; ++i)
  {
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaFree(sendbuff[i]));
    CUDACHECK(cudaFree(recvbuff[i]));
  }

  if(NULL != sendbuff_H)
  {
    free(sendbuff_H);
    sendbuff_H = NULL;
  }
  if(NULL != recvbuff_H)
  {
    free(recvbuff_H);
    recvbuff_H = NULL;
  }
  //finalizing NCCL
  for(int i = 0; i < nDev; ++i)
  {
    ncclCommDestroy(comms[i]);
  }

  printf("=====================end==============\n");
  return 0;
}

static void broadcast_root(int nDev, int root, int nCount, int size, int* sendbuff_H, int* recvbuff_H, int** sendbuff, int** recvbuff, cudaStream_t* stream, ncclComm_t *comms)
{
  NCCLCHECK(ncclGroupStart());
  for (int i = 0; i < nDev; ++i)
  {

    CUDACHECK(cudaSetDevice(i));

    printf("[%s +%d %s] call ncclBroadcast root=%d\n", __FILE__, __LINE__, __func__, root);
    NCCLCHECK(ncclBroadcast((const void*)sendbuff[i], (void*)recvbuff[i], nCount, ncclInt, root, comms[i], stream[i]));
    printf("[%s +%d %s] call ncclBroadcast end\n", __FILE__, __LINE__, __func__);
  }

  printf("[%s +%d %s] call ncclGroupEnd start\n", __FILE__, __LINE__, __func__);
  NCCLCHECK(ncclGroupEnd());
  printf("[%s +%d %s] call ncclGroupEnd end\n", __FILE__, __LINE__, __func__);

  //synchronizing on CUDA streams to wait for completion of NCCL operation
  for (int i = 0; i < nDev; ++i)
  {
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaStreamSynchronize(stream[i]));
  }

  for (int i = 0; i < nDev; ++i)
  {
    printf("[%s +%d %s] copy recvbuff[%d] to recvbuff_H by cudaMemcpy\n", __FILE__, __LINE__, __func__, i);
    CUDACHECK(cudaMemcpy(recvbuff_H, recvbuff[i], size, cudaMemcpyDeviceToHost));
    for (int j = 0; j < nCount; ++j)
    {
      //printf("[%s +%d %s] [%d] *(recvbuff_H + j=%d)=%d\n", __FILE__, __LINE__, __func__, i, j, *(recvbuff_H + j));
    }
  }
}